from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings

from langchain_text_splitters import CharacterTextSplitter

docs = [
    Document(
        page_content="LangChain 是一个让大模型可以使用外部数据和工具的框架。",
        # metadata 可以包含文档的额外信息，如来源、类别、创建时间等
        metadata={"source": "langchain_intro", "category": "technical"}
    ),
    Document(
        page_content="Retriever 是 LangChain 中用于从数据库中检索相关信息的组件。",
        metadata={"source": "retriever_guide", "category": "technical"}
    ),
    Document(
        page_content="Milvus 是一个高性能的向量数据库，用于存储和检索向量化的数据。",
        metadata={"source": "milvus_docs", "category": "database"}
    )
]

# 切分文档
splitter = CharacterTextSplitter(chunk_size=30, chunk_overlap=10)
split_docs = splitter.split_documents(docs)

embeddings = HuggingFaceEmbeddings(model_name="D:/models/BAAI_bge-small-zh-v1.5", model_kwargs={'device': 'cpu'})

# 使用现有的 VectorStoreRetriever，但自定义参数
vectorstore = Chroma.from_documents(split_docs, embeddings)

# 自定义检索器参数
custom_retriever = vectorstore.as_retriever(
    search_type="similarity",       # 搜索类型：相似度/MMR
    search_kwargs={
        "k": 5,                     # 返回前5个最相关文档
        #"score_threshold": 0.8,     # 只返回相似度>0.8的文档
        "filter": {"category": "technical"}  # 只检索技术类文档
    }
)

print("\n🔍 测试检索器功能:")
test_queries = [
    "什么是 LangChain？",
    "Retriever 有什么作用？",
    "向量数据库是什么？"
]

for query in test_queries:
    print(f"\n查询: '{query}'")

    # 使用检索器查找相关文档
    relevant_docs = custom_retriever.get_relevant_documents(query)

    print(f"找到 {len(relevant_docs)} 个相关文档:")
    for i, doc in enumerate(relevant_docs):
        print(f"  {i + 1}. {doc.page_content}")
        print(f"     相似度分数: {doc.metadata.get('score', 'N/A')}")
        print(f"     元数据: {doc.metadata}")